import os
import zipfile
import requests
from tqdm import tqdm


def support_filename_encode(zip_file: zipfile.ZipFile, encoding='utf-8'):
    """
    zipfile读取zip包中的中文会使用`cp437`方式，导致出现乱码
    :param zip_file: 读取文件得到的ZipFile对象
    :param encoding: 希望改变后的编码方式，默认是 'utf-8'
    :return: 修改NameToInfo结构后的ZipFile对象
    """
    name_to_info = zip_file.NameToInfo
    # 先copy一个对象
    for name, info in name_to_info.copy().items():
        real_name = name.encode('cp437').decode(encoding)
        if real_name != name:
            # 修改文件名
            info.filename = real_name
            # 删除，并且重新设置文明名到文件的映射关系
            del name_to_info[name]
            name_to_info[real_name] = info

    return zip_file


def download_dataset(file_path, file_url, chunk_size=1024):
    """
    通过网络下载所需要的数据集攻击
    :param file_path: 文件保存路径
    :param file_url: 数据集下载url
    :param chunk_size: 分块下载时分块大小
    """
    if os.path.exists(file_path):
        os.remove(file_path)
    req = requests.get(file_url, stream=True)
    file_size = int(req.headers.get('Content-Length')) / chunk_size
    with open(file_path, 'wb') as fp:
        for item in tqdm(iterable=req.iter_content(chunk_size),
                         desc='download dataset',
                         total=file_size):
            fp.write(item)

    print(f'Download success! Dataset path is: {file_path}')
